package main

import (
	"net/http"

	"golang.org/x/time/rate"
)

var r = rate.NewLimiter(1, 5)

func main() {
	// 速度为1, 容量为5
	// r := rate.NewLimiter(1, 5)
	// ctx := context.Background()

	// for {
	// 	err := r.WaitN(ctx, 2)
	// 	if err != nil {
	// 		log.Panicln(err)
	// 	}
	// 	fmt.Println(time.Now().Format("2006-01-02 15:04:05"))
	// 	// time.Sleep(1 * time.Second)
	// }

	// r := rate.NewLimiter(1, 5)
	// for {
	// 	if r.AllowN(time.Now(), 2) {
	// 		fmt.Println(time.Now().Format("2006-01-02 15:04:05"))
	// 	} else {
	// 		log.Println("to many request!")
	// 	}
	// 	time.Sleep(1 * time.Second)
	// }

	// http server
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(rw http.ResponseWriter, r *http.Request) {
		rw.Write([]byte("ok!!!!"))
	})

	http.ListenAndServe(":8080", myLimit(mux))
}

func myLimit(next http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, request *http.Request) {
		if !r.Allow() {
			http.Error(rw, "too many requests", http.StatusTooManyRequests)
			return
		}
		next.ServeHTTP(rw, request)
	})
}
